#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
蓝奏云下载链接与验证码提取工具
主模块文件，提供命令行入口点
"""

import argparse
import sys
import traceback
import time
import json
import os
from modules.config_manager import ConfigManager
from modules.browser_manager import BrowserManager
from modules.login_handler import LoginHandler
from modules.search_handler import SearchHandler
from modules.link_extractor import LinkExtractor
from modules.download_handler import DownloadHandler
from modules.utils import print_banner, print_section, confirm_action, setup_logging

def main():
    """主函数，处理命令行参数并启动自动化流程"""
    # 显示欢迎横幅
    print_banner()
    
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='蓝奏云下载链接与验证码提取工具')
    parser.add_argument('-s', '--search', help='直接指定搜索关键词')
    parser.add_argument('-c', '--comment', default='感谢分享', help='设置默认评论内容')
    parser.add_argument('-u', '--url', help='直接处理指定URL，跳过搜索步骤')
    parser.add_argument('--no-headless', action='store_true', help='不使用无头模式（显示浏览器界面）')
    parser.add_argument('--debug', action='store_true', help='启用调试模式，显示更多信息')
    parser.add_argument('--auto', action='store_true', help='自动下载所有链接，不提示确认', default=True)
    parser.add_argument('-n', '--number', type=int, help='指定要处理的搜索结果数量，默认处理所有结果')
    parser.add_argument('--select', action='store_true', help='手动选择要处理的搜索结果，而不是自动处理')
    parser.add_argument('--ask', action='store_true', help='下载前询问确认，默认为自动下载')
    parser.add_argument('--loop', action='store_true', help='开启循环模式，下载完成后可继续输入新关键词', default=False)
    parser.add_argument('--batch', help='批量模式：指定包含搜索内容和数量的JSON文件路径')
    parser.add_argument('--batch-create', action='store_true', help='创建批量配置模板文件')
    args = parser.parse_args()
    
    # 处理批量配置模板创建请求
    if args.batch_create:
        create_batch_config_template()
        return 0
    
    # 加载配置
    config_manager = ConfigManager()
    config = config_manager.load_config()
    
    # 初始化浏览器
    browser_manager = BrowserManager(headless=False, debug=args.debug)
    
    # 加载批量配置（如果有）
    batch_configs = []
    if args.batch:
        try:
            if not os.path.exists(args.batch):
                print(f"错误：批量配置文件 '{args.batch}' 不存在")
                return 1
            
            with open(args.batch, 'r', encoding='utf-8') as f:
                batch_data = json.load(f)
                
                if not isinstance(batch_data, list):
                    print("错误：批量配置文件格式不正确，应为列表")
                    return 1
                
                for item in batch_data:
                    if isinstance(item, dict) and 'keyword' in item:
                        num_to_process = item.get('number', None)
                        batch_configs.append((item['keyword'], num_to_process))
                    else:
                        print(f"警告：跳过无效的配置项 {item}")
            
            print(f"已加载 {len(batch_configs)} 个批量搜索配置")
        except json.JSONDecodeError:
            print(f"错误：'{args.batch}' 不是有效的JSON文件")
            return 1
        except Exception as e:
            print(f"读取批量配置文件时出错: {str(e)}")
            return 1
    
    try:
        # 启动浏览器
        driver = browser_manager.initialize_browser()
        if not driver:
            print("无法初始化浏览器，程序退出")
            return 1
        
        # 初始化各个处理模块
        login_handler = LoginHandler(driver, config)
        search_handler = SearchHandler(driver)
        link_extractor = LinkExtractor(driver)
        download_handler = DownloadHandler(driver)
        
        # 处理登录
        if config['require_login']:
            print_section("登录系统")
            login_handler.login()
        
        # 是否需要询问下载确认
        auto_download = not args.ask
        
        # 循环执行直到用户选择退出
        continue_loop = True
        
        # 如果有批量配置，优先处理
        if batch_configs:
            print_section("批量处理搜索任务")
            total_tasks = len(batch_configs)
            
            for idx, (keyword, number) in enumerate(batch_configs, 1):
                print(f"\n【{idx}/{total_tasks}】正在处理关键词: {keyword}")
                
                # 执行搜索
                results = search_handler.perform_search(keyword)
                
                # 处理搜索结果
                if results and len(results) > 0:
                    # 确定要处理的结果数量
                    if number:
                        num_to_process = min(number, len(results))
                    else:
                        num_to_process = len(results)
                    
                    print(f"共找到 {len(results)} 个结果，将处理前 {num_to_process} 个")
                    
                    # 选择结果
                    selected_results = []
                    for i in range(num_to_process):
                        selected_results.append(results[i])
                        print(f"{i+1}. {results[i]['title']}")
                    
                    # 处理选中的结果
                    if selected_results:
                        process_search_results(
                            driver, 
                            selected_results, 
                            link_extractor, 
                            download_handler, 
                            args.comment, 
                            auto_download=auto_download
                        )
                    else:
                        print("未选择任何搜索结果进行处理")
                else:
                    print(f"关键词 '{keyword}' 未找到任何搜索结果")
                
                # 完成一个批量任务后，如果不是最后一个，回到主页准备下一次搜索
                if idx < total_tasks:
                    print("\n正在准备下一个搜索任务...")
                    driver.get("https://www.hifini.com/")
                    time.sleep(2)
            
            print_section("批量任务完成")
            print(f"已完成所有 {total_tasks} 个批量搜索任务")
            
            # 批量任务完成后，如果设置了循环模式，可以继续进入交互模式
            if not args.loop:
                continue_loop = False
        
        while continue_loop:
            # 如果直接指定了URL，则跳过搜索步骤
            if args.url:
                print_section(f"直接处理指定URL: {args.url}")
                driver.get(args.url)
                process_single_page(driver, link_extractor, download_handler, args.comment, auto_download=auto_download)
                
                # URL模式下只执行一次，除非开启了循环模式
                if not args.loop:
                    continue_loop = False
                else:
                    args.url = None  # 清除URL，下次进入搜索模式
            else:
                # 搜索和处理结果
                print_section("搜索内容")
                search_text = args.search
                if not search_text:
                    print("\n请输入要搜索的内容 (输入'exit'或'quit'退出): ")
                    search_text = input().strip()
                    
                    # 检查是否退出
                    if search_text.lower() in ['exit', 'quit']:
                        print("用户选择退出程序")
                        continue_loop = False
                        break
                else:
                    # 如果是命令行参数传入的搜索词，只执行一次，除非开启了循环模式
                    if not args.loop:
                        args.search = None  # 清除搜索词，避免重复使用
                
                # 执行搜索
                if search_text:
                    results = search_handler.perform_search(search_text)
                    
                    # 处理搜索结果
                    if results and len(results) > 0:
                        # 询问用户希望处理多少个链接
                        if args.number:
                            # 如果命令行指定了链接数量，直接使用
                            num_to_process = min(args.number, len(results))
                            print(f"\n将处理前 {num_to_process} 个搜索结果")
                        else:
                            # 否则询问用户
                            print(f"\n共找到 {len(results)} 个搜索结果，请输入想要处理的数量 (默认全部): ")
                            num_input = input().strip()
                            if num_input and num_input.isdigit():
                                num_to_process = min(int(num_input), len(results))
                            else:
                                num_to_process = len(results)
                            print(f"将处理前 {num_to_process} 个搜索结果")
                        
                        # 根据用户输入选择结果
                        selected_results = []
                        for i in range(num_to_process):
                            selected_results.append(results[i])
                            print(f"{i+1}. {results[i]['title']}")
                        
                        # 处理选中的结果
                        if selected_results:
                            process_search_results(
                                driver, 
                                selected_results, 
                                link_extractor, 
                                download_handler, 
                                args.comment, 
                                auto_download=auto_download
                            )
                        else:
                            print("未选择任何搜索结果进行处理")
                    else:
                        print("未找到任何搜索结果")
            
            # 完成一轮操作后，询问是否继续新的搜索
            if continue_loop:
                print_section("操作完成")
                print("当前任务已处理完毕")
                if args.loop or confirm_action("是否开始新的搜索?"):
                    # 如果处于循环模式或用户确认，继续新的搜索
                    print("\n开始新的搜索...")
                    # 确保回到搜索页或主页
                    driver.get("https://www.hifini.com/")
                else:
                    continue_loop = False
        
        print_section("完全退出")
        print("所有任务已处理完毕")
        
        # 询问是否关闭浏览器
        if confirm_action("是否关闭浏览器?"):
            browser_manager.close_browser()
        else:
            print("浏览器将保持打开状态，您可以继续浏览")
            input("\n按回车键退出程序（浏览器将保持打开状态）...")
        
        return 0
        
    except KeyboardInterrupt:
        print("\n\n操作已被用户中断。")
        print("感谢使用本工具。再见!")
        return 0
    except Exception as e:
        print(f"\n发生错误: {str(e)}")
        print("请尝试重新运行脚本或联系开发者寻求帮助。")
        traceback.print_exc()
        return 1
    finally:
        # 确保浏览器正常关闭（如果需要）
        try:
            if 'browser_manager' in locals() and 'driver' in locals():
                browser_manager.close_browser()
        except:
            pass

def create_batch_config_template():
    """创建批量配置模板文件"""
    template = [
        {
            "keyword": "例如：周杰伦",
            "number": 3
        },
        {
            "keyword": "例如：林俊杰",
            "number": 2
        },
        {
            "keyword": "例如：陈奕迅"
            # 不指定number则处理所有结果
        }
    ]
    
    filename = "batch_config_template.json"
    with open(filename, 'w', encoding='utf-8') as f:
        json.dump(template, f, ensure_ascii=False, indent=4)
    
    print(f"批量配置模板已创建: {filename}")
    print("请编辑此文件，填入您要搜索的内容和对应要处理的数量")
    print("然后使用 --batch 参数指定此文件路径运行脚本")
    print("例如: python lanzou_downloader.py --batch batch_config.json")

def process_search_results(driver, selected_results, link_extractor, download_handler, comment_text, auto_download=False):
    """处理选中的搜索结果"""
    processed_count = 0
    
    for i, result in enumerate(selected_results):
        print_section(f"处理结果 {i+1}/{len(selected_results)}: {result['title']}")
        
        # 在新标签页打开链接
        driver.execute_script(f"window.open('{result['url']}');")
        
        # 切换到新标签页
        driver.switch_to.window(driver.window_handles[-1])
        
        # 处理单个页面
        success = process_single_page(driver, link_extractor, download_handler, comment_text, auto_download)
        if success:
            processed_count += 1
        
        # 如果不是最后一个结果，关闭当前标签页并返回到搜索结果页
        if i < len(selected_results) - 1:
            # 自动处理下一个，不需要询问
            print(f"已处理 {i+1}/{len(selected_results)} 个链接，继续处理下一个...")
            
            driver.close()
            driver.switch_to.window(driver.window_handles[0])
    
    if processed_count > 0:
        print(f"\n成功处理了 {processed_count} 个蓝奏云下载链接")
    else:
        print("\n未找到或处理任何蓝奏云下载链接")

def process_single_page(driver, link_extractor, download_handler, comment_text, auto_download=True):
    """处理单个页面，返回是否成功处理"""
    # 提取蓝奏云下载链接和验证码
    links_data = link_extractor.extract_multiple_lanzou_links()
    
    if links_data and len(links_data) > 0:
        print(f"\n找到 {len(links_data)} 个蓝奏云链接:")
        
        for i, (link, pwd) in enumerate(links_data, 1):
            print(f"{i}. 链接: {link}, 密码: {pwd}")
        
        # 处理找到的链接
        success_count = 0
        
        # 直接下载所有链接，不再询问
        for i, (link, pwd) in enumerate(links_data, 1):
            print(f"\n处理链接 {i}/{len(links_data)}: {link}")
            print(f"正在处理下载链接 {i}/{len(links_data)}...")
            # 使用auto_process=True强制自动处理下载过程
            try:
                if download_handler.process_download_link(link, pwd, auto_process=True):
                    success_count += 1
                    print(f"链接 {i} 下载成功")
                else:
                    print(f"链接 {i} 下载失败")
            except Exception as e:
                print(f"处理链接 {i} 时出错: {str(e)}")
                # 继续处理下一个链接
                continue
        
        print(f"\n成功处理了 {success_count}/{len(links_data)} 个蓝奏云链接")
        return success_count > 0
    else:
        print("\n未找到蓝奏云链接")
        
        # 询问是否尝试发表评论
        try_comment = auto_download or confirm_action("是否尝试回帖评论以显示下载链接?")
        
        if try_comment:
            if link_extractor.post_comment(comment_text):
                # 重新检查下载链接
                links_data = link_extractor.extract_multiple_lanzou_links()
                
                if links_data and len(links_data) > 0:
                    print(f"\n评论后找到 {len(links_data)} 个蓝奏云链接:")
                    
                    for i, (link, pwd) in enumerate(links_data, 1):
                        print(f"{i}. 链接: {link}, 密码: {pwd}")
                    
                    # 处理所有找到的链接
                    success_count = 0
                    
                    # 直接下载所有链接，不再询问
                    for i, (link, pwd) in enumerate(links_data, 1):
                        print(f"\n处理链接 {i}/{len(links_data)}: {link}")
                        print(f"正在处理下载链接 {i}/{len(links_data)}...")
                        try:
                            if download_handler.process_download_link(link, pwd, auto_process=True):
                                success_count += 1
                                print(f"链接 {i} 下载成功")
                            else:
                                print(f"链接 {i} 下载失败")
                        except Exception as e:
                            print(f"处理链接 {i} 时出错: {str(e)}")
                            # 继续处理下一个链接
                            continue
                    
                    print(f"\n成功处理了 {success_count}/{len(links_data)} 个蓝奏云链接")
                    return success_count > 0
                else:
                    print("\n评论后仍未找到蓝奏云链接")
        
        # 询问是否查看页面上的其他链接
        check_links = auto_download or confirm_action("是否查看页面上的所有链接?")
        
        if check_links:
            specific_url = link_extractor.check_all_links(auto_select=auto_download)
            if specific_url:
                print(f"正在打开选择的链接: {specific_url}")
                driver.get(specific_url)
                
                # 在新页面再次尝试提取蓝奏云链接
                links_data = link_extractor.extract_multiple_lanzou_links()
                
                if links_data and len(links_data) > 0:
                    print(f"\n在新页面找到 {len(links_data)} 个蓝奏云链接:")
                    
                    for i, (link, pwd) in enumerate(links_data, 1):
                        print(f"{i}. 链接: {link}, 密码: {pwd}")
                    
                    # 处理所有找到的链接
                    success_count = 0
                    
                    # 直接下载所有链接
                    for i, (link, pwd) in enumerate(links_data, 1):
                        print(f"\n处理链接 {i}/{len(links_data)}: {link}")
                        print(f"正在处理下载链接 {i}/{len(links_data)}...")
                        try:
                            if download_handler.process_download_link(link, pwd, auto_process=True):
                                success_count += 1
                                print(f"链接 {i} 下载成功")
                            else:
                                print(f"链接 {i} 下载失败")
                        except Exception as e:
                            print(f"处理链接 {i} 时出错: {str(e)}")
                            # 继续处理下一个链接
                            continue
                    
                    print(f"\n成功处理了 {success_count}/{len(links_data)} 个蓝奏云链接")
                    return success_count > 0
        
        # 询问是否手动输入HTML提取链接
        if not auto_download and confirm_action("是否手动输入HTML内容提取链接?"):
            links_data = link_extractor.manual_html_extract_links()
            
            if links_data and len(links_data) > 0:
                print(f"\n从HTML内容中找到 {len(links_data)} 个蓝奏云链接:")
                
                for i, (link, pwd) in enumerate(links_data, 1):
                    print(f"{i}. 链接: {link}, 密码: {pwd}")
                
                # 处理所有找到的链接
                success_count = 0
                
                # 直接下载所有链接
                for i, (link, pwd) in enumerate(links_data, 1):
                    print(f"\n处理链接 {i}/{len(links_data)}: {link}")
                    print(f"正在处理下载链接 {i}/{len(links_data)}...")
                    try:
                        if download_handler.process_download_link(link, pwd, auto_process=True):
                            success_count += 1
                            print(f"链接 {i} 下载成功")
                        else:
                            print(f"链接 {i} 下载失败")
                    except Exception as e:
                        print(f"处理链接 {i} 时出错: {str(e)}")
                        # 继续处理下一个链接
                        continue
                
                print(f"\n成功处理了 {success_count}/{len(links_data)} 个蓝奇云链接")
                return success_count > 0
    
    return False

if __name__ == "__main__":
    sys.exit(main()) 